#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : partitioner.py
@Project  : pipecoco
@Date     : 2021/8/13
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

from pipecoco.utils import *
import numpy as np

def layers_area_partition(size, block_num, layers, output_shapes):
    """
    融合网络各层区域划分
    :param size: 样本输入大小
    :param block_num: 分块数量
    :param layers: 顺序为由上至下的各网络层
    :param output_shapes: 各层的特征图输出维度
    :return:
        area_map: 每层要计算的区域
        pre_area_map: 每层要计算的区域映射到前一层的区域
    """

    # 计算融合网络最顶层划分方式，保证计算均衡，将多余的计算单元分配给排序较后的区域
    area_map = top_layer_area_partition_by_proportion(size, block_num)

    # 用于存放计算area_map中区域单元要用到的上一层对应区域
    pre_area_map = []

    for i in range(block_num):
        for j in range(block_num):
            border_x, border_y = area_map[i * block_num + j][-1]
            pre_area_map.append([])
            x = 0
            # 由上至下推导各层分块区域
            for _, layer in enumerate(layers):
                if isinstance(layer,Bypass):
                    continue

                kernel_size, stride, p = layer.kernel_size, layer.stride, layer.padding
                # 找到上一层网络对应的区域并加入pre_layer_area的映射中
                border_x, border_y = pre_layer_area(border_x, border_y, kernel_size, stride)
                pre_area_map[i * block_num + j].append([border_x, border_y])

                # 去除区域中超出上一层特征图边界的部分和重复的部分
                output_shape = output_shapes[x + 1]
                if border_x[1] > border_x[0]:
                    border_x = [max(0, border_x[0] - p[0]), min(border_x[1] - p[1], output_shape[1]-2*p[1])]
                    if i > 0:
                        border_x = [max(area_map[(i - 1) * block_num + j][x + 1][0][1], border_x[0]), border_x[1]]

                if border_y[1] > border_y[0]:
                    border_y = [max(0, border_y[0] - p[0]), min(border_y[1] - p[1], output_shape[2]-2*p[1])]
                    if j > 0:
                        border_y = [max(area_map[i * block_num + j - 1][x + 1][1][1], border_y[0]), border_y[1]]

                # 将要计算的区域加入特征图中
                area_map[i * block_num + j].append([border_x, border_y])
                x+=1

    # 将区域分块重新按照由下之上网络层顺序排布
    for i in range(block_num):
        for j in range(block_num):
            area_map[i * block_num + j] = area_map[i * block_num + j][::-1]
            pre_area_map[i * block_num + j] = pre_area_map[i * block_num + j][::-1]
    return np.array(area_map, dtype=np.int64), np.array(pre_area_map, dtype=np.int64)


def top_layer_area_partition_by_diff(size, block_num):
    """
    根据数值差的融合网络最顶层划分方法
    由于映射到下层网络区域存在重叠部分，通过倾向于减小先计算的block区域的方式实现均衡负载
    :param size: 顶层特征图大小
    :param block_num: 分块数量
    :return:
    """
    area_map = []
    height = size[0]
    width = size[1]

    rest_h = height % block_num
    rest_w = width % block_num

    avg_h = height // block_num
    avg_w = width // block_num

    # 先计算的区域长宽较小
    h_units = [avg_h for _ in range(block_num - rest_h)] + [avg_h + 1 for _ in range(rest_h)]
    w_units = [avg_w for _ in range(block_num - rest_w)] + [avg_w + 1 for _ in range(rest_w)]

    if rest_h == 0:
        h_units[0] = h_units[0] - 1
        h_units[-1] = h_units[-1] + 1

    if rest_w == 0:
        w_units[0] = w_units[0] - 1
        w_units[-1] = w_units[-1] + 1

    # 得到最顶层的区域划分方式
    h = 0
    for i, h_unit in enumerate(h_units):
        border_x = (h, h + h_unit)
        w = 0
        for j, w_unit in enumerate(w_units):
            border_y = (w, w + w_unit)
            w += w_unit
            area_map.append([[border_x, border_y]])
        h += h_unit
    return area_map


def top_layer_area_partition_by_proportion(size, block_num):
    """
    根据比例设计的融合网络最顶层划分方法
    由于映射到下层网络区域存在重叠部分，通过倾向于减小先计算的block区域的方式实现均衡负载
    :param size: 顶层特征图大小
    :param block_num: 分块数量
    :return:
    """
    # divide_proportion = {
    #     1:[],
    #     2:[1/3],
    #     3:[1/6,1/3],
    #     4:[1/8,1/4,1/4]
    # }
    divide_proportion = {
        1:[],
        2:[2/5],
        3:[2/9,1/3],
        4:[3/16,1/4,1/4],
        5:[4/30,6/30,6/30,6/30]
    }
    area_map = []
    height = size[0]
    width = size[1]

    scheme = divide_proportion[block_num]
    # 先计算的区域长宽较小
    h_units = [int(height*prop) for prop in scheme]
    h_units.append(height-sum(h_units))
    w_units = [int(width*prop) for prop in scheme]
    w_units.append(width - sum(w_units))

    # 得到最顶层的区域划分方式
    h = 0
    for i, h_unit in enumerate(h_units):
        border_x = (h, h + h_unit)
        w = 0
        for j, w_unit in enumerate(w_units):
            border_y = (w, w + w_unit)
            w += w_unit
            area_map.append([[border_x, border_y]])
        h += h_unit
    return area_map



def cal_shape(input_shape, fused_layers):
    """
    根据输入维度和网络结构计算各层输出维度
    :param input_shape: 输入维度
    :param fused_layers: 各层网络
    :return:
    """
    c, h, w = input_shape
    output_shape = []
    for j, layer in enumerate(fused_layers):
        if isinstance(layer, Bypass):
            continue
        h += 2 * layer.padding[0]
        w += 2 * layer.padding[1]
        output_shape.append([c, h, w])
        # 计算下一层输出的大小
        h, w = next_layer_size(h, w, layer.kernel_size, layer.stride)

        if hasattr(layer, 'out_channels'):
            c = layer.out_channels

    output_shape.append([c, h, w])
    return output_shape


def next_layer_size(h, w, kernel_size = 1, stride = 1, p = (0,0)):
    """
    计算当前层区域映射到下一层的区域
    """

    h = (h + 2 * p[0] - kernel_size) // stride + 1
    w = (w + 2 * p[1] - kernel_size) // stride + 1
    return h, w


def pre_layer_area(h, w, kernel_size, stribe):
    """
    计算当前层区域映射到的上一层区域
    :param h: 长
    :param w: 宽
    :param kernel_size: 核尺寸
    :param stribe: 步长
    :return:
    """
    x1, x2 = h
    y1, y2 = w
    if not isinstance(kernel_size, tuple):
        kernel_size = (kernel_size, kernel_size)

    if not isinstance(stribe, tuple):
        stribe = (stribe, stribe)
    if x2>x1:
        border_x = [x1 * stribe[0], (x2 - 1) * stribe[0] + kernel_size[0]]
    else:
        border_x = [x1 * stribe[0], x1 * stribe[0]]
    if y2>y1:
        border_y = [y1 * stribe[1], (y2 - 1) * stribe[1] + kernel_size[1]]
    else:
        border_y = [y1 * stribe[0], y1 * stribe[0]]
    return border_x, border_y



def cal_workload(area_map,fuse_layers):
    """
    计算量统计，卷积网络中计算量主要来自MAC计算量
    计算方法：C_in * K * K * H_out * W_out + C_in * C_out * H_out * W_out
    :param area_map: 各分块计算区域
    :param fuse_layers: 融合网络层
    :return:
    """
    s = 0
    for i,layer in enumerate(fuse_layers):
        if not isinstance(layer,pipecoco_layer) or not isinstance(layer.layer[0],nn.Conv2d):
            continue
        area = area_map[i+1]

        HW = (area[0][1]-area[0][0])*(area[1][1]-area[1][0])
        C_in = layer.in_channels
        C_out = layer.out_channels
        if isinstance(layer.kernel_size,tuple):
            K = layer.kernel_size[0]
        else:
            K = layer.kernel_size
        # print(C_in*K*K*HW + C_in*C_out*HW)
        s += C_in*K*K*HW + C_in*C_out*HW

    return s
